from __future__ import division, print_function

import matplotlib.pyplot as plt
import matplotlib
import numpy as np
import cv2 as cv
import os
import pdb
plt.rcParams['image.cmap'] = 'gist_earth'
np.random.seed(98765)

import sys
sys.path.append("..")
from tf_unet import image_gen
from tf_unet import unet
from tf_unet import util
from tf_unet.image_util import ImageDataProvider, ImageDataSingle
import tensorflow as tf

# cmap = {"background":[0,0,0], "lane_line":[128,0,0], "tree":[0,128,0], "tank":[128,128,0], "building":[0,0,128], "plane":[128,0,128], "wzw":[0,128,128]}
label_cmap = {0:[0,0,0], 1:[128,0,0], 2:[0,128,0], 3:[128,128,0], 4:[0,0,128], 5:[128,0,128], 6:[0,128,128]}
label_cmap_list = np.array(list(label_cmap.values()))

# data_provider = ImageDataSingle("./data/cj_cut.png", "./data/cj_cut_gt.png")
# data_provider = ImageDataSingle("./data/cj_test1.png", "./data/cj_test1_gt.png")
data_provider = ImageDataSingle("./data/cj_0402_1420/cj_right_all_20200402_1420.png", "./data/cj_0402_1420/cj_right_all_20200402_1420_gt.png")

output_path = "./log_l3_best"
n_class = 7
class_weights = np.ones(n_class)*5
class_weights[0] = 0.5
net = unet.Unet(data_provider.channels, n_class, layers=3,
                features_root=32, cost_kwargs={"class_weights": class_weights})

def predict_panoroma(path, data_provider):
    pass

def stack_imgs(imgs, num_row, num_col):
    '''
    concatenate image slices to a panoroma,
    imgs: image slices should be sorted by row first
    '''
    imgs_row = []
    for i in range(num_row):
        imgs_row.append(np.concatenate(imgs[i*num_col:(i+1)*num_col], axis=1))
    img_stack = np.concatenate(imgs_row, axis=0)
    return img_stack.astype(np.uint8)

path=tf.train.latest_checkpoint(output_path) 


images = data_provider.images_origin
masks = data_provider.masks_origin
prediction = net.predict(path, images)

gts = np.argmax(masks, axis=-1)
images_crop = util.crop_to_shape(images, prediction.shape)
gts_crop = util.crop_to_shape(gts, prediction.shape)

preds_rgb = []
preds_gt = []
for i, pred in enumerate(prediction):
    h, w = pred.shape[0:2]
    label = pred.argmax(axis=-1)
    preds_gt.append(label)
    label = label.reshape([-1])
    label_rgb = np.array(label_cmap_list[label])
    label_rgb = label_rgb.reshape((h, w, 3))
    preds_rgb.append(label_rgb)

img = stack_imgs(images_crop, data_provider.num_row, data_provider.num_col)
img_gt = stack_imgs(gts_crop, data_provider.num_row, data_provider.num_col)
pred_rgb = stack_imgs(preds_rgb, data_provider.num_row, data_provider.num_col)
pred_gt = stack_imgs(preds_gt, data_provider.num_row, data_provider.num_col)
cv.imwrite("pred_rgb.png", pred_rgb)
cv.imwrite("img.png", img)
img_add = cv.addWeighted(img, 0.8, pred_rgb, 0.2, 0)
cv.imwrite("img_add.png", img_add)
uniq = img_gt == pred_gt
print("accuracy: %.4f" %(uniq.sum()/np.product(img_gt.shape)))

# for temp use, because label order is different
img_gt_tmp = np.zeros_like(img_gt)
indexs = dict(zip([1,2,3,4,5], [4, 5, 3, 2, 6]))
for i in range(1, 6):
    mask = img_gt == i
    img_gt_tmp[mask] = indexs[i]
uniq = img_gt_tmp == pred_gt
print("accuracy: %.4f" %(uniq.sum()/np.product(img_gt.shape)))

for i, c in enumerate(label_cmap_list):
    if i == 0: continue
    pred_rgb[100:150, i*100:(i*100+50)] = c
    cv.putText(pred_rgb, str(i), (i*100+50, 100), 1, 2, (255, 255, 255))
cv.imwrite("test.png", pred_rgb)


# img = cv.imread("./data/cj_right_all.png", -1)
# data = []
# nx = 572
# for i in range(0, img.shape[1]-nx, nx):
#     img_tmp = img[:nx, i:i+nx]
#     data.append(img_tmp)
# data = np.array(data)

# prediction = net.predict(path, data)
# data_crop = util.crop_to_shape(data, prediction.shape)
# for i, img in enumerate(data_crop):
#     pred_img = np.zeros_like(img)
#     for l in range(1,data_provider.n_class):
#         mask = prediction[i, :, :, l] > 0.1
#         color = label_cmap[l]
#         pred_img[mask] = color
#     img_tmp = cv.addWeighted(img, 0.8, pred_img, 0.2, 0)
#     cv.imwrite("pred_%d.png" %i, img_tmp)
#     cv.imwrite("pred_%d.png" %i, pred_img)
